
use core::{
    cell::UnsafeCell,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

use super::waker_set::WakerSet;
use crate::executor::RESUME;
use core::sync::atomic::Ordering;
use alloc::collections::LinkedList;

pub struct Channel<T> {
    buffer: UnsafeCell<LinkedList<T>>,
    recv_wakers: WakerSet,
}

impl<T> Channel<T> {
    /// Creates a new fixed capacity channel
    pub fn new() -> Self {
        Self {
            buffer: UnsafeCell::new(LinkedList::new()),
            recv_wakers: WakerSet::new(),
        }
    }

    /// Sends a message into the channel
    pub async fn send(&self, val: T) {
        self.try_send(val).unwrap_err();
    }

    /// Receives a message from the channel
    pub async fn recv(&self) -> T {
        struct Recv<'a, T> {
            channel: &'a Channel<T>,
            opt_key: Option<usize>,
        }

        impl<T> Future for Recv<'_, T> {
            type Output = T;

            fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
                // If the current task is in the set, remove it.
                if let Some(key) = self.opt_key.take() {
                    self.channel.recv_wakers.remove(key);
                }

                // Try receiving a message.
                if let Some(msg) = self.channel.try_recv() {
                    Poll::Ready(msg)
                } else {
                    // Insert this receive operation.
                    self.opt_key = Some(self.channel.recv_wakers.insert(cx));
                    Poll::Pending
                }
            }
        }

        Recv {
            channel: self,
            opt_key: None,
        }
        .await
    }

    pub fn try_recv(&self) -> Option<T> {
        unsafe {
            let buf = &mut *(self.buffer.get());
            buf.pop_front()
        }
    }

    pub fn try_send(&self, val: T) -> Result<(), T> {
        unsafe {
            let buf = &mut *(self.buffer.get());
            buf.push_back(val);
            self.recv_wakers.notify_one();
            RESUME.store(true, Ordering::Release);
            Ok(())
        }
    }
}
